import os
import io
import time
import pywt
import base64
from tempo import *
from tempo0 import *  # 含有setParameters等函数
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


def add_noise_to_trajectory(trajectory, noise_std):  # 加噪函数
    """
    向航迹数据中添加高斯噪声。

    参数:
    trajectory (numpy.ndarray): 形状为 (N, D) 的航迹数据，其中 N 是点的数量，D 是每个点的维度。
    noise_std (float): 噪声的标准差。

    返回:
    numpy.ndarray: 添加了噪声的航迹数据。
    """
    # 生成高斯噪声
    noise = np.random.normal(0, noise_std, size=trajectory.shape)
    # 添加噪声
    noisy_trajectory = trajectory + noise
    return noisy_trajectory


def wavelet_denoise(data, wavelet='haar', level=4, mode='symmetric'):
    """
    使用小波变换对数据进行去噪。
    sym4  haar  db4
    参数:
    data (numpy.ndarray): 待去噪的数据。
    wavelet (str): 使用的小波基函数，默认为 'db4'。
    level (int): 小波分解的层数，默认为 1。
    mode (str): 边界处理模式，默认为 'symmetric'。

    返回:
    numpy.ndarray: 去噪后的数据。
    """

    def sure_shrink(coeffs, noise_std):
        """
        自适应阈值选择：SUREShrink
        :param coeffs: 小波变换得到的系数列表
        :param noise_std: 噪声的标准差估计
        :return: 优化后的系数列表
        """
        def soft_threshold(coeff, threshold):
            return np.sign(coeff) * np.maximum(np.abs(coeff) - threshold, 0)
        
        # 初始化最优阈值和风险
        best_risk = np.inf
        best_threshold = None
        
        # 对每个尺度的系数进行自适应阈值选择
        for detail_level in range(1, len(coeffs)):
            # 初始阈值
            initial_threshold = noise_std * np.sqrt(2 * np.log(len(coeffs[detail_level])))
            
            # 在一定范围内搜索最佳阈值
            thresholds = np.linspace(0, 2 * initial_threshold, 100)
            risks = []
            
            for threshold in thresholds:
                # 应用软阈值处理
                coeffs_shrinked = soft_threshold(coeffs[detail_level], threshold)
                
                # 计算 SURE 风险估计
                risk = np.sum(coeffs_shrinked ** 2) + 2 * threshold * np.sum(np.abs(coeffs_shrinked)) - noise_std ** 2 * len(coeffs[detail_level])
                risks.append(risk)
            
            # 寻找最小风险对应的阈值
            min_risk_idx = np.argmin(risks)
            best_threshold = thresholds[min_risk_idx]
            
            # 更新系数
            coeffs[detail_level] = soft_threshold(coeffs[detail_level], best_threshold)

        return coeffs

    # level = pywt.dwt_max_level(len(data), pywt.Wavelet(wavelet).dec_len)

    coeffs = pywt.wavedec(data, wavelet, level=level,mode=mode)

    # 计算噪声标准差的估计
    mad = lambda x: np.median(np.abs(x - np.median(x))) / 0.6745
    sigma = mad(coeffs[-level])

    noise_std = sigma
    # 对每个尺度的系数进行自适应阈值处理
    coeffs_denoised = sure_shrink(coeffs, noise_std)
    # 进行小波逆变换
    data_denoised = pywt.waverec(coeffs_denoised, wavelet,mode=mode)
    # 确保输出长度与输入长度一致
    if len(data_denoised) != len(data):
        diff = len(data) - len(data_denoised)
        if diff > 0:
            data_denoised = np.pad(data_denoised, (0, diff), mode='edge')
        elif diff < 0:
            data_denoised = data_denoised[:len(data)]    
    return data_denoised


def plot_n_trajectory(groups, labels, colors, show=False, save_path=None):
    """Plot trajectories and return three separate base64 encoded images."""
    image_base64_list = []

    # Create three separate figures instead of subplots
    for i in range(3):
        fig = plt.figure(figsize=(8, 6))
        if i == 0:
            ax = fig.add_subplot(111, projection='3d')
        else:
            ax = fig.add_subplot(111)

        for group, label, color in zip(groups, labels, colors):
            # Convert timestamp to datetime
            times_dt = pd.to_datetime(group['time'], unit='s')
            times_num = mdates.date2num(times_dt)

            # Extract values from the group
            latitudes = group.lat.values
            longitudes = group.lon.values
            geoAltitudes = group.geoaltitude.values
            times = group.time.values

            if i == 0:  # 3D trajectory
                ax.scatter3D(longitudes, latitudes, geoAltitudes, s=2, color=color, label=label)
                ax.set_title('3D trajectory')
                ax.set_xlabel('Longitude')
                ax.set_ylabel('Latitude')
                ax.set_zlabel('GeoAltitude')
            elif i == 1:  # 2D scatter plot
                ax.scatter(longitudes, latitudes, s=2, color=color, label=label)
                ax.set_title('2D scatter plot')
                ax.set_xlabel('Longitude')
                ax.set_ylabel('Latitude')
            else:  # Altitude plot
                ax.plot(times_num, geoAltitudes, color=color, label=label)
                ax.set_title('Altitude plot')
                ax.set_xlabel('Time')
                ax.set_ylabel('GeoAltitude')
                ax.set_xticks(ax.get_xticks()[::2])  # keep every other tick
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))

        ax.legend()

        if show:
            plt.tight_layout()
            plt.show()

            # Save individual plot if save_path is provided
        if save_path is not None:
            fig.savefig(f"{save_path}_{i}.png")

            # Convert to base64
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
        image_base64_list.append(image_base64)
        plt.close()

        # Return three separate base64 encoded images
    return image_base64_list


def set_and_load_parameters_1(parameters_1):
    """设置和加载功能一的参数"""
    setParameters(parameters_1)
    return {
    "trajdata_path": parameters_1["trajdata_path"],
    }


def load_and_visualize_trajectory(parameters):
    """功能一：导入数据"""
    # 设置和加载功能一的参数
    params = set_and_load_parameters_1(parameters)
    trajdata_path = params["trajdata_path"]
    # 加载航迹数据
    org_data = np.load(trajdata_path, allow_pickle=True)
    np_data = np.array(org_data, dtype=float)
    # 归一化
    cols_to_normalize = [1, 2, 3]
    data_to_normalize = np_data[:, cols_to_normalize]
    min_vals = data_to_normalize.min(axis=0)
    max_vals = data_to_normalize.max(axis=0)
    ranges = max_vals - min_vals
    ranges[ranges == 0] = 1
    data_normalized = (data_to_normalize - min_vals) / ranges
    data_normalized_full = np_data.copy()
    data_normalized_full[:, cols_to_normalize] = data_normalized
    # 可视化原始数据
    columns = ['time', 'lat', 'lon', 'geoaltitude']
    traj_df = pd.DataFrame(data_normalized_full, columns=columns)
    image_base64_original = plot_n_trajectory([traj_df], ['Original'], ['red'])
    # 输出反归一化参数min_val、 ranges
    denormalized_para = [min_vals, ranges]
    # 输出
    setShowData({"image_base64": image_base64_original})

    return {
        'traj_data': data_normalized_full,
        'org_data': np_data,
        'image_base64_original': image_base64_original,
        'denormalized_para': denormalized_para
        }


def set_and_load_parameters_2(parameters_2):
    """设置和加载功能二的参数"""
    setParameters(parameters_2)
    return {
    "noise_std": parameters_2["noise_std"],
    }


def add_noise_and_visualize_trajectories(parameters_2, data):
    """功能二：加噪"""
    # 设置和加载功能二的参数
    params = set_and_load_parameters_2(parameters_2)
    noise_std = params["noise_std"]
    # 创建一个空列表来存储非时间列的数据
    noisy_data = []
    image_base64_original_noised_all = []
    columns = ['time', 'lat', 'lon', 'geoaltitude']
    for i in range(len(data)):
        non_time_data = data[i][:, 1:4]
        noisy_trajectory0 = add_noise_to_trajectory(non_time_data, noise_std)
        time_column = data[i][:, 0]
        combined_data = np.column_stack((time_column, noisy_trajectory0))
        noisy_data.append(combined_data)
        noise_traj_df = pd.DataFrame(combined_data, columns=columns)
        trajectory = data[i]
        traj_df = pd.DataFrame(trajectory, columns=columns)
        image_base64_original_noised = plot_n_trajectory([traj_df, noise_traj_df], ['Original', 'Noise'],
                                                         ['red', 'green'])
        setShowData({"image_base64": image_base64_original_noised})
        image_base64_original_noised_all.append(image_base64_original_noised)

    return {'noisy_data': noisy_data, 'image_base64_original_noised': image_base64_original_noised_all}


def set_and_load_parameters_3(parameters_3):
    """设置和加载功能三的参数"""
    setParameters(parameters_3)
    return {"Lat_Lon_wavelet_type": parameters_3["Lat_Lon_wavelet_type"],
            "Lat_Lon_level": parameters_3["Lat_Lon_level"],
            "Lat_Lon_mode": parameters_3["Lat_Lon_mode"],
            "Alt_wavelet_type": parameters_3["Alt_wavelet_type"],
            "Alt_level": parameters_3["Alt_level"],
            "Alt_mode": parameters_3["Alt_mode"],
            "denoised_save_path": parameters_3["denoised_save_path"]}


def calculate_dataset_errors(tra_data_all, denoised_data):
    """
    计算整个数据集的去噪误差

    参数:
    tra_data_all: list[2D:np.array, float64], 原始航迹数据集
    denoised_data: list[2D:np.array, float64], 去噪后的航迹数据集

    返回:
    dict: 包含各种误差指标的字典
    """
    start_time = time.time()

    # 存储所有航迹的误差
    all_errors = []
    # 遍历每条航迹
    for original_traj, denoised_traj in zip(tra_data_all, denoised_data):
        # 确保两个轨迹数据长度相同
        assert len(original_traj) == len(denoised_traj), "轨迹长度不匹配"

        # 计算当前航迹的所有点的误差
        traj_errors = np.sqrt(
            (original_traj[:, 1] - denoised_traj[:, 1]) ** 2 +  # 经度误差
            (original_traj[:, 2] - denoised_traj[:, 2]) ** 2  # 纬度误差
        )

        all_errors.extend(traj_errors)

        # 转换为numpy数组便于计算
    all_errors = np.array(all_errors)

    # 计算各种误差指标
    mse = np.mean(all_errors ** 2)
    rmse = np.sqrt(mse)
    avg_error = np.mean(all_errors)
    max_error = np.max(all_errors)
    median_error = np.median(all_errors)
    std_error = np.std(all_errors)

    # 计算运行时间
    execution_time = time.time() - start_time

    # 构建评估结果字典
    eval_result = {
        "运行时间(s)": execution_time,
        "均方误差(MSE)": mse,
        "均方根误差(RMSE)": rmse,
        "平均误差": avg_error,
        "最大误差": max_error,
        "中位数误差": median_error,
        "误差标准差": std_error, "数据点总数": len(all_errors), "轨迹条数": len(tra_data_all)
    }

    # 输出结果
    setShowData({"eval_result": eval_result})

    return eval_result


def denoise_trajectories_with_wavelet(parameters_3, lines, tra_data_all, noisy_data,denormalized_paras):
    """功能三：去噪"""
    start_time = time.time()
    # 设置和加载功能三的参数
    params = set_and_load_parameters_3(parameters_3)
    lat_lon_wavelet_type = params["Lat_Lon_wavelet_type"]
    lat_lon_level = params["Lat_Lon_level"]
    lat_lon_mode = params["Lat_Lon_mode"]
    alt_wavelet_type = params["Alt_wavelet_type"]
    alt_level = params["Alt_level"]
    alt_mode = params["Alt_mode"]
    denoised_save_path = params["denoised_save_path"]
    if not os.path.exists(denoised_save_path):
        os.makedirs(denoised_save_path)
    # 去噪处理
    denoised_data = []
    image_base64_all = []
    denoiesd_denormalized_datas = []
    columns = ['time', 'lat', 'lon', 'geoaltitude']
    for i in range(len(noisy_data)):
        denoise_data = np.zeros_like(noisy_data[i])
        denoise_data[:, 0] = noisy_data[i][:, 0]
        for col in range(noisy_data[i].shape[1]):
            if col == 0:
                continue
            elif col == 3:
                denoise_data[:, col] = wavelet_denoise(noisy_data[i][:, col],
                                                       alt_wavelet_type,
                                                       alt_level,
                                                       alt_mode)
            else:
                denoise_data[:, col] = wavelet_denoise(noisy_data[i][:, col], lat_lon_wavelet_type, lat_lon_level,
                                                       lat_lon_mode)
        denoised_data.append(denoise_data)
        trajectory = tra_data_all[i]
        traj_df = pd.DataFrame(trajectory, columns=columns)
        noisy_trajectory = noisy_data[i]
        noise_traj_df = pd.DataFrame(noisy_trajectory, columns=columns)
        denoised_df = pd.DataFrame(denoise_data, columns=columns)
        image_base64 = plot_n_trajectory([traj_df, noise_traj_df, denoised_df],
                                         ['Original', 'Noise', 'Denoised'],
                                         ['red', 'green', 'blue'])
        setShowData({"image_base64": image_base64})
        image_base64_all.append(image_base64)
        file_name = lines[i].split('\\')[-1]  # 获取文件名
        save_path = f"{denoised_save_path}\\{file_name[:-4]}_denoised.npy"
        #反归一化
        denoiesd_denormalized_data  =denoise_data[:, [1, 2, 3]]
        denoiesd_denormalized_data = denoiesd_denormalized_data * denormalized_paras[i][1] + denormalized_paras[i][0]#denormalized_paras第一列为min_vals,第二列为ranges
        denoiesd_denormalized_data_full = denoise_data.copy()
        denoiesd_denormalized_data_full[:, [1, 2, 3]] = denoiesd_denormalized_data
        denoiesd_denormalized_datas.append(denoiesd_denormalized_data_full)
        np.save(save_path, denoiesd_denormalized_data_full)
    # 计算误差
    eval_result = calculate_dataset_errors(tra_data_all, denoised_data)
    print(eval_result)

    return {
        'denoised_data': denoised_data,
        'denoiesd_denormalized_datas': denoiesd_denormalized_datas,
        'image_base64_all': image_base64_all,
        'eval_result': eval_result
        }


if __name__ == "__main__":
    start_time = time.time()
    '''功能一：加载航迹数据'''
    # parameters_1 = {"trajdata_path": r'D:\pythonProject\去噪20241031\原始数据\ab197e_18885.npy'}
    # setParameters(parameters_1)
    # load_and_visualize_result = load_and_visualize_trajectory(parameters_1)
    # load_and_visualize_result = {'traj_data': data_normalized_full, 'org_data': org_data,
    #                              'image_base64_original': image_base64_original, 'denormalized_para':denormalized_para}
    # org_data                 2D:np.array，float64                          加载的原始航迹数据，用于在地图上展示
    # traj_data:               2D:np.array，float64                          归一化的航迹数据，功能二加噪模块的输入
    # image_base64_original:   list[string1(base64编码),string2,string3]     所有轨迹的原始图像
    # denormalized_para        list[np.array,np.array]                       用于反变换

    """加载所有数据"""
    # 获取文件路径
    folder_path = r'D:\pythonProject\去噪20241031\原始数据'
    file_name = os.listdir(folder_path)
    lines = []
    for alphabet in file_name:
        lines.append(folder_path + '\\' + alphabet)
    # 循环调用
    tra_data_all = []
    denormalized_paras = []
    for file_path in lines:
        # 构建参数字典
        parameters = {
            "trajdata_path": file_path
        }
        result = load_and_visualize_trajectory(parameters)
        tra_data = result['traj_data']
        tra_data_all.append(tra_data)
        denormalized_paras.append(result['denormalized_para'])                      # 保存各原始航迹归一化参数
    '''功能二：加噪'''
    parameters_2 = {"noise_std": 0.01}                                              # float类型，加噪标准差，高斯白噪声
    setParameters(parameters_2)
    add_noise_and_visual_result = add_noise_and_visualize_trajectories(parameters_2, tra_data_all)
    # add_noise_and_visual_result = {'noisy_data': noisy_data, 'image_base64_original_noised': image_base64_original_noised_all}
    # noisy_data:                   list[2D:np.array，float64]                         加噪声的航迹数据，功能三去噪模块的输入
    # image_base64_original_noised: list[list[string1(base64编码),string2,string3] ]   所有轨迹的加噪图像

    noisy_data = add_noise_and_visual_result['noisy_data']

    '''功能三：去噪'''
    parameters_3 = {"Lat_Lon_wavelet_type": 'haar',                                 # string类型，小波变换基函数的类型，做下拉菜单，可选sym4  haar  db4
                    "Lat_Lon_level": 4,                                             # int类型，小波分解层数，一般四层左右效果较好
                    "Lat_Lon_mode": 'symmetric',                                    # string类型，小波变换边界处理函数，做下拉菜单，可选symmetric
                                                                                    # periodic reflect constant zero smooth antireflect
                    "Alt_wavelet_type": 'db4',                                      # string类型，小波变换基函数的类型，做下拉菜单，可选sym4  haar  db4
                    "Alt_level": 2,                                                 # int类型，小波分解层数，一般四层左右效果较好
                    "Alt_mode": 'symmetric',                                        # string类型，小波变换边界处理函数，做下拉菜单，可选symmetric
                                                                                    # periodic reflect constant zero smooth antireflect
                    "denoised_save_path": r'D:\pythonProject\去噪20241031\去噪航迹'} # string类型，去噪航迹保存的文件夹路径
    setParameters(parameters_3)
    denoise_trajectories_with_wavelet_result = denoise_trajectories_with_wavelet(parameters_3, lines, tra_data_all, noisy_data,denormalized_paras)
    execution_time = time.time() - start_time
    print(execution_time)
    # denoise_trajectories_with_wavelet_result = {'denoised_data': denoised_data, 'image_base64_all': image_base64_all,
    #                                             'eval_result': eval_result}
    # denoised_data:                list[np.array]，float64                           去去噪的结果
    # denoiesd_denormalized_datas   list[np.array]                                    去噪反变换的结果
    # eval_result:                  dict,包含的键:  ['运行时间', '均方误差', '平均误差', '最大误差', '中位数误差']，值numpy.float64
    # image_base64_all：            list[list[string1(base64编码),string2,string3] ]   所有轨迹的去噪图像
